#include "ANN_Header.h"

#include "WMatrix.h"

namespace ann{

void TickFunctor( ANN_Node* node )
{
	node->Update();
}

void NodeInitFunctor( ANN_Node* node )
{
    node->SetValue( NODE_INIT_VALUE );
}

void TubeInitFunctor( Tube tube )
{
    tube.SetValue(  TUBE_INIT_VALUE );
    tube.SetWeight( WEIT_INIT_VALUE );
}

BPNet::BPNet(   const unsigned int _num_in  /*= 2*/,
                const unsigned int _num_bp  /*= 2*/,
                const unsigned int _num_out /*= 1*/,
                const double       _alpha   /*= 0.6*/,
                const double       _p_criEk /*= 0.008*/,
                const unsigned int p_itrLmt /*= 8000*/
                ):
    m_isInited      ( false ),
    m_isConverged   ( false ),
    m_cnt_trainData ( 0     ),
    m_itrCnt        ( 0     ),
    m_criEk         ( _p_criEk ),
    m_criEk_double  ( 2.0 * _p_criEk ),
    m_itrLmt        ( p_itrLmt ),
    m_alpha         ( _alpha ),
    m_cnt_inNode    ( _num_in ),
    m_cnt_bpNode    ( _num_bp ),
    m_cnt_outNode   ( _num_out ),
    m_cnt_tube      ( _num_in * _num_bp + _num_bp * _num_out ),
    m_cnt_valveTube ( _num_out + _num_bp ),
    m_ek            ( NULL ),
    m_cnt_layer     ( 2 ) //temp for a fixed layer count
{
    m_inNodeList    = new InputNode[m_cnt_inNode];
    m_bpNodeList    = new BPNode[m_cnt_bpNode];
    m_outNodeList   = new OutputNode[m_cnt_outNode];
    m_tubeList      = new Tube[m_cnt_tube];
    m_valveTubeList = new Tube[m_cnt_valveTube];
    
    //wmat
    m_wMat.clear();
    m_bVec.clear();
    m_dim.clear();
    
    m_dim.push_back(_num_out);
    m_dim.push_back(_num_bp);
    m_dim.push_back(_num_in);

    //TODO: check whether -1 works ok
    for( uiVItr itr = m_dim.begin(); itr < m_dim.end()-1; ++itr )
    {
        WMatrix* wm = new WMatrix(*itr, *(itr+1));
        dVec     b(*itr, 0.0);

        m_wMat.push_back(wm);
        m_bVec.push_back(b);
    }
}

bool BPNet::Init()
{
    //the init value of weight must be set to different value
    srand( (long)time(NULL) );

    //std::for_each( m_inNodeList,    m_inNodeList + m_cnt_inNode,    NodeInitFunctor );
    //std::for_each( m_bpNodeList,    m_bpNodeList + m_cnt_bpNode,    NodeInitFunctor );
    //std::for_each( m_outNodeList,   m_outNodeList+ m_cnt_outNode,   NodeInitFunctor );

	for ( unsigned int i = 0; i < m_cnt_inNode; ++i )
	{
		NodeInitFunctor( &m_inNodeList[i] );
	}
	for ( unsigned int i = 0; i < m_cnt_bpNode; ++i )
	{
		NodeInitFunctor( &m_bpNodeList[i] );
	}
	for ( unsigned int i = 0; i < m_cnt_outNode; ++i )
	{
		NodeInitFunctor( &m_outNodeList[i] );
	}

    //init valve
    for ( unsigned int i = 0; i < m_cnt_valveTube; ++i )
    {
        m_valveTubeList[i].SetValue(VALV_INIT_VALUE);
        m_valveTubeList[i].SetWeight(VALV_WEIT_INIT);
    }

    //std::for_each( m_tubeList,      m_tubeList + m_cnt_tube,        TubeInitFunctor );
	for ( unsigned int i = 0; i < m_cnt_tube; ++i )
	{
		m_tubeList[i].SetValue( TUBE_INIT_VALUE );
		m_tubeList[i].SetWeight( WEIT_INIT_VALUE);
	}

    //+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    //manually input init value, to ensure converge
    m_valveTubeList[0] = -0.0109;
    m_valveTubeList[1] = -0.0703;
    m_valveTubeList[2] = -0.0939;
    m_tubeList[0] = 0.0801;
    m_tubeList[1] = -0.0605;
    m_tubeList[2] = 0.0543;
    m_tubeList[3] = -0.0291;
    m_tubeList[4] = 0.0579;
    m_tubeList[5] = 0.0999;
    //+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    //add inNode and tube connection
    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
    {
        for( unsigned int j = 0; j < m_cnt_bpNode; ++j )
        {
            m_outNodeList[i].AddInTube( &m_tubeList[GetTubeIDX(i, j, UNRELATED_IDX)] );
            m_tubeList[GetTubeIDX(i, j, UNRELATED_IDX)].SetDestNode( &m_outNodeList[i] );
        }
        m_outNodeList[i].AddInTube( &m_valveTubeList[i] );
    }

    for( unsigned int i = 0; i < m_cnt_bpNode; ++i )
    {
        for( unsigned int j = 0; j < m_cnt_inNode; ++j )
        {
            m_bpNodeList[i].AddInTube( &m_tubeList[GetTubeIDX(UNRELATED_IDX, i, j)] );
            m_tubeList[GetTubeIDX(UNRELATED_IDX, i, j)].SetDestNode( &m_bpNodeList[i] );
        }
        m_bpNodeList[i].AddInTube( &m_valveTubeList[i + m_cnt_outNode] );
    }

    //add outNode and tube connection
    for( unsigned int i = 0; i < m_cnt_inNode; ++i )
    {
        for( unsigned int j = 0; j < m_cnt_bpNode; ++j )
        {
            m_inNodeList[i].AddOutTube( &m_tubeList[GetTubeIDX(UNRELATED_IDX, j, i)] );
            m_tubeList[GetTubeIDX(UNRELATED_IDX, j, i)].SetSrcNode( &m_inNodeList[i] );
        }
    }

    for( unsigned int i = 0; i < m_cnt_bpNode; ++i )
    {
        for( unsigned int j = 0; j < m_cnt_outNode; ++j )
        {
            m_bpNodeList[i].AddOutTube( &m_tubeList[GetTubeIDX(j, i, UNRELATED_IDX)] );
            m_tubeList[GetTubeIDX(j, i, UNRELATED_IDX)].SetSrcNode( &m_bpNodeList[i] );
        }
    }


    //wmat, init
    //TODO: file reading parse, temp hard code
    double** b = new double*[2];
    b[0] = new double[1];
    b[1] = new double[2];
    double** m = new double*[2];
    m[0] = new double[2];
    m[1] = new double[4];
    b[0][0] = -0.0109;
    b[1][0] = 0.0703;
    b[1][1] = -0.0939;
    m[0][0] = 0.0801;
    m[0][1] = -0.0605;
    m[1][0] = 0.0543;
    m[1][1] = -0.0291;
    m[1][2] = 0.0579;
    m[1][3] = 0.0999;
    for( unsigned int i = 0; i < m_cnt_layer; ++i )
    {
        for( unsigned int j = 0; j < m_dim[i]; ++j )
        {
            m_bVec[i][j] = b[i][j];

            for( unsigned int k = 0; k < m_dim[i+1]; ++k )
            {
                (*m_wMat[i])(j,k) = m[j][k];
            }
        }
    }

    m_isInited = true;
	
	return m_isInited;
}

bool BPNet::Reset()
{
    return Init();
}

bool BPNet::InputData( const std::vector<unsigned int> &_p_inData )
{
    if ( _p_inData.size() == m_cnt_inNode )
    {
        for( unsigned int i = 0; i < m_cnt_inNode; ++i )
        {
            m_inNodeList[i].InputData( _p_inData[i] );
        }
        return true;
    }

    return false;
}

dVec BPNet::GetResult() const
{
    dVec retVec;

    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
    {
        retVec.push_back( m_outNodeList[i].GetValue() );
    }

	return *(new dVec(retVec));
}

bool BPNet::PushTrainData( const TrainData _tData )
{
    if ( _tData._s_inData.size() == m_cnt_inNode && _tData._s_result.size() == m_cnt_outNode )
    {
        m_trainData.push_back(_tData);
		++m_cnt_trainData;
        return true;
    }

    return false;
}

bool BPNet::Train()
{
    m_isConverged = false;

	unsigned int itrCnt = 0;

    m_ek = new double[m_cnt_trainData];
    std::fill( m_ek, m_ek + m_cnt_trainData, 0.0 );

	DEBUG("----------------------Start Training--------------------------\n");
    for( ; !this->IsConverged() && itrCnt<m_itrLmt; this->TestConverge() )
    {
		for ( unsigned int i = 0; i < m_cnt_trainData; ++i )
		{
			this->InputData( m_trainData[i]._s_inData );
			this->Tick();
			this->Update_W(i);
            this->CalEk(i);
		}
		++itrCnt;
    }
    DEBUG("Train times = %d, Converged at: \n", --itrCnt);
    //TODO: output the weight matrix
	DEBUG("----------------------End Training--------------------------\n");

    if ( NULL != m_ek )
    {
        delete [] m_ek;
        m_ek = NULL;
    }

    return this->IsConverged();
}

unsigned int BPNet::GetTubeIDX( const unsigned int _p_outNode_id,
								const unsigned int _p_bpNode_id,
								const unsigned int _p_inNode_id )
{
    assert( (_p_outNode_id == -1 || _p_inNode_id == -1) && _p_bpNode_id != -1 );

    unsigned int idx = 0;
    if ( _p_outNode_id == -1 )
    {
        idx = m_cnt_outNode * m_cnt_bpNode + _p_bpNode_id * m_cnt_inNode + _p_inNode_id;
    }
    else
    {
        idx = _p_outNode_id * m_cnt_outNode + _p_bpNode_id;
    }

    return idx;
}

void BPNet::Tick()
{
    //std::for_each( m_inNodeList,    m_inNodeList + m_cnt_inNode,    TickFunctor );
    //std::for_each( m_bpNodeList,    m_bpNodeList + m_cnt_bpNode,    TickFunctor );
    //std::for_each( m_outNodeList,   m_outNodeList+ m_cnt_outNode,   TickFunctor );

	for ( unsigned int i = 0; i < m_cnt_inNode; ++i )	m_inNodeList[i].Update();
	for ( unsigned int i = 0; i < m_cnt_bpNode; ++i )	m_bpNodeList[i].Update();
	for ( unsigned int i = 0; i < m_cnt_outNode; ++i )	m_outNodeList[i].Update();
}


void Tick_M()
{
    //TODO:
    //finds some big defeat in weight matrix design
    //so, have to let it for further further investigating
}

void BPNet::PropogateBack()
{
}

void BPNet::Update_W( const unsigned int _p_trainIdx )
{
    double *_deltaW= new double[m_cnt_tube];
    double *_deltaV= new double[m_cnt_valveTube];

    //the deltaW calculation part
    double *o2 = new double[m_cnt_outNode];
    double *o1 = new double[m_cnt_bpNode];
    double *o0 = new double[m_cnt_inNode];
    double *dp = new double[m_cnt_outNode];
    double *dp2= new double[m_cnt_outNode];
    double *dp1= new double[m_cnt_bpNode];

    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
        o2[i] = m_outNodeList[i].GetValue();

    for( unsigned int i = 0; i < m_cnt_bpNode; ++i )
        o1[i] = m_bpNodeList[i].GetValue();

    for( unsigned int i = 0; i < m_cnt_inNode; ++i )
        o0[i] = fabs( m_inNodeList[i].GetValue() - EPSILON );

    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
        dp[i] = fabs( static_cast<double>(m_trainData[_p_trainIdx]._s_result[i]) - EPSILON );

    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
        dp2[i] = 2.0 * ( dp[i] - o2[i]) * o2[i] * ( 1.0 - o2[i] );

    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
    {
        for( unsigned int j = 0; j < m_cnt_bpNode; ++j )
            _deltaW[ GetTubeIDX( i, j, UNRELATED_IDX ) ] = m_alpha * dp2[i] * o1[j];

        _deltaV[i] = m_alpha * dp2[i] * VALV_INIT_VALUE;
    }

    for( unsigned int i = 0; i < m_cnt_bpNode; ++i )
    {
        double sigmaDW = 0.0;
        for( unsigned int j = 0; j < m_cnt_outNode; ++j )
            sigmaDW += dp2[j] * (m_tubeList[ GetTubeIDX( j, i, UNRELATED_IDX ) ].GetWeight() + _deltaW[GetTubeIDX( j, i, UNRELATED_IDX )]);
        dp1[i] = sigmaDW * o1[i] * ( 1 - o1[i] );
    }

    for( unsigned int i = 0; i < m_cnt_bpNode; ++i )
    {
        for( unsigned int j = 0; j < m_cnt_inNode; ++j )
            _deltaW[ GetTubeIDX( UNRELATED_IDX, i, j ) ] = m_alpha * dp1[i] * o0[j];

        _deltaV[i + m_cnt_outNode] = m_alpha * dp1[i] * VALV_INIT_VALUE;
    }
    //end of the deltaW calculation part

    for( unsigned int i = 0; i < m_cnt_tube; ++i )
        m_tubeList[i].SetWeight( m_tubeList[i].GetWeight() + _deltaW[i] );

    for( unsigned int i = 0; i < m_cnt_valveTube; ++i )
        m_valveTubeList[i].SetWeight( m_valveTubeList[i].GetWeight() + _deltaV[i] );
}

void BPNet::CalEk( const unsigned int p_trainIdx )
{ 
    if ( m_isConverged || p_trainIdx>=m_trainData.size() )
    {
        return;
	}

    double* o = new double[m_cnt_outNode];
    double* r = new double[m_cnt_outNode];

    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
    {
        o[i] = m_outNodeList[i].GetValue();
    }

    uVector& expVec = m_trainData[p_trainIdx]._s_result;
    std::copy( expVec.begin(), expVec.end(), r );
	
    double sumEk = 0.0;
    double rdo   = 0.0;
    for( unsigned int i = 0; i < m_cnt_outNode; ++i )
    {
        rdo   = r[i] - o[i];
        sumEk = rdo * rdo;
    }

    m_ek[p_trainIdx] = sumEk;
}

bool BPNet::TestConverge()
{
    double maxEk = *std::max_element( m_ek, m_ek + m_cnt_trainData );
    if ( maxEk < m_criEk_double )
    {
        DEBUG("Converged, the maxEk is %f\n", maxEk*0.5 );
        m_isConverged = true;
    }

	return m_isConverged;
}

bool BPNet::IsConverged() const
{
	return m_isConverged;
}

BPNet::~BPNet(void)
{
    delete [] m_inNodeList;
    delete [] m_bpNodeList;
    delete [] m_outNodeList;
    delete [] m_tubeList;
}

}
